package com.idea.dsl.client;

import com.fasterxml.jackson.databind.ObjectMapper;
import com.idea.dsl.config.DslProperties;
import com.idea.dsl.model.ModelRequest;
import com.idea.dsl.model.ModelResponse;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import okhttp3.*;
import org.apache.commons.lang3.StringUtils;
import org.springframework.stereotype.Component;

import java.io.IOException;
import java.util.HashMap;
import java.util.Map;

/**
 * 本地模型客户端实现（如Ollama）
 */
@Slf4j
@Component
@RequiredArgsConstructor
public class LocalModelClient implements ModelClient {
    
    private final OkHttpClient httpClient;
    private final ObjectMapper objectMapper;
    private final DslProperties dslProperties;
    
    private static final String CLIENT_TYPE = "local";
    private static final MediaType JSON = MediaType.get("application/json; charset=utf-8");
    
    @Override
    public ModelResponse chat(ModelRequest request) {
        try {
            DslProperties.ModelConfig config = dslProperties.getModels().get(CLIENT_TYPE);
            if (config == null) {
                throw new RuntimeException("本地模型配置未找到");
            }
            
            // 构建本地模型请求格式（以Ollama为例）
            Map<String, Object> localRequest = buildLocalRequest(request, config);
            String requestBody = objectMapper.writeValueAsString(localRequest);
            log.debug("本地模型请求: {}", requestBody);
            
            // 构建HTTP请求
            Request httpRequest = new Request.Builder()
                    .url(config.getBaseUrl() + "/generate")
                    .header("Content-Type", "application/json")
                    .post(RequestBody.create(requestBody, JSON))
                    .build();
            
            // 发送请求
            try (Response response = httpClient.newCall(httpRequest).execute()) {
                if (!response.isSuccessful()) {
                    throw new RuntimeException("本地模型请求失败: " + response.code() + " " + response.message());
                }
                
                String responseBody = response.body().string();
                log.debug("本地模型响应: {}", responseBody);
                
                return convertFromLocalFormat(responseBody);
            }
            
        } catch (IOException e) {
            log.error("本地模型请求异常", e);
            throw new RuntimeException("本地模型请求异常: " + e.getMessage(), e);
        }
    }
    
    private Map<String, Object> buildLocalRequest(ModelRequest request, DslProperties.ModelConfig config) {
        Map<String, Object> localRequest = new HashMap<>();
        localRequest.put("model", config.getModel());
        
        // 合并所有消息内容
        StringBuilder prompt = new StringBuilder();
        if (request.getMessages() != null) {
            for (ModelRequest.Message message : request.getMessages()) {
                if (StringUtils.isNotBlank(message.getContent())) {
                    prompt.append(message.getContent()).append("\n");
                }
            }
        }
        
        localRequest.put("prompt", prompt.toString().trim());
        localRequest.put("stream", false);
        
        Map<String, Object> options = new HashMap<>();
        if (request.getMaxTokens() != null) {
            options.put("num_predict", request.getMaxTokens());
        }
        if (request.getTemperature() != null) {
            options.put("temperature", request.getTemperature());
        }
        localRequest.put("options", options);
        
        return localRequest;
    }
    
    private ModelResponse convertFromLocalFormat(String responseBody) throws IOException {
        Map<String, Object> localResponse = objectMapper.readValue(responseBody, Map.class);
        
        String content = (String) localResponse.get("response");
        if (content == null) {
            content = "";
        }
        
        ModelRequest.Message message = ModelRequest.Message.builder()
                .role("assistant")
                .content(content)
                .build();
        
        ModelResponse.Choice choice = ModelResponse.Choice.builder()
                .index(0)
                .message(message)
                .finishReason("stop")
                .build();
        
        return ModelResponse.builder()
                .id("local-" + System.currentTimeMillis())
                .model((String) localResponse.get("model"))
                .choices(java.util.Arrays.asList(choice))
                .created(System.currentTimeMillis() / 1000)
                .build();
    }
    
    @Override
    public String getClientType() {
        return CLIENT_TYPE;
    }
    
    @Override
    public boolean isAvailable() {
        DslProperties.ModelConfig config = dslProperties.getModels().get(CLIENT_TYPE);
        return config != null && StringUtils.isNotBlank(config.getBaseUrl());
    }
}
